#!/usr/bin/env python3
"""
Generate PBPK Parameters from Population Data

This script extracts physiological parameters from the population folder
and converts them into the format required by individual_pbpk_modeling.py
and population_pbpk_modeling.py

Author: BlackCat@CPPO
Version: 1.0.0
"""

import pandas as pd
import numpy as np
import sys
import os
import json

# Add population path to system path
project_root = os.path.dirname(os.path.abspath(__file__))
sys.path.insert(0, project_root)

class PopulationParameterGenerator:
    """Generate PBPK parameters from population data"""
    
    def __init__(self, population_csv='population/china_population.csv', 
                 total_params_csv='population/total_population_parameters.csv'):
        self.population_csv = os.path.join(project_root, population_csv)
        self.total_params_csv = os.path.join(project_root, total_params_csv)
        self.population_df = None
        self.total_params_df = None
        self.load_data()
    
    def load_data(self):
        """Load population data"""
        try:
            # Try to load existing population data
            if os.path.exists(self.total_params_csv):
                self.total_params_df = pd.read_csv(self.total_params_csv)
                print(f"Loaded total population parameters from {self.total_params_csv}")
            
            if os.path.exists(self.population_csv):
                self.population_df = pd.read_csv(self.population_csv)
                print(f"Loaded population data from {self.population_csv}")
            else:
                print("Population data not found, using synthetic data...")
                self.generate_synthetic_data()
                
        except Exception as e:
            print(f"Error loading data: {e}")
            self.generate_synthetic_data()
    
    def generate_synthetic_data(self):
        """Generate synthetic population data based on Chinese population"""
        np.random.seed(42)
        n = 100
        
        # Generate basic demographics
        sexes = np.random.choice(['M', 'F'], n, p=[0.5, 0.5])
        ages = np.random.uniform(20, 80, n)
        heights_male = np.random.normal(172, 6, sum(sexes == 'M'))
        heights_female = np.random.normal(160, 5, sum(sexes == 'F'))
        heights = np.concatenate([heights_male, heights_female])
        
        # Calculate weights based on height and age
        bmis = np.random.normal(23, 3, n)  # BMI distribution
        weights = bmis * (heights / 100) ** 2
        
        # Calculate BSA using DuBois formula
        bsas = 0.007184 * (weights ** 0.425) * (heights ** 0.725)
        
        self.population_df = pd.DataFrame({
            'id': range(1, n+1),
            'sex': sexes,
            'age': ages,
            'height': heights,
            'weight': weights,
            'bsa': bsas
        })
        
        # Generate organ parameters
        self.total_params_df = self.population_df.copy()
        
        # Add liver parameters
        self.total_params_df['Liver_weight'] = 0.025 * self.total_params_df['weight']
        self.total_params_df['Liver_volume'] = self.total_params_df['Liver_weight'] / 1.08  # density 1.08 g/mL
        self.total_params_df['Liver_blood_flow'] = 1.35 * self.total_params_df['bsa'] * 60  # mL/min to L/h
        self.total_params_df['Liver_CYP3A4_activity'] = np.random.lognormal(0, 0.3, n)
        
        # Add kidney parameters
        self.total_params_df['Kidney_weight'] = 0.004 * self.total_params_df['weight']
        self.total_params_df['Kidney_volume'] = self.total_params_df['Kidney_weight'] / 1.05
        self.total_params_df['Kidney_blood_flow'] = 1.25 * self.total_params_df['bsa'] * 60
        self.total_params_df['Kidney_GFR'] = np.random.normal(120, 20, n)
        
        # Add heart parameters
        self.total_params_df['Heart_weight'] = 0.0045 * self.total_params_df['weight']
        self.total_params_df['Heart_volume'] = self.total_params_df['Heart_weight'] / 1.05
        self.total_params_df['Heart_blood_flow'] = 5.0 * self.total_params_df['bsa'] * 60
        
        # Add lung parameters
        self.total_params_df['Lung_weight'] = 0.01 * self.total_params_df['weight']
        self.total_params_df['Lung_volume'] = self.total_params_df['Lung_weight'] / 0.3
        self.total_params_df['Lung_blood_flow'] = 5.5 * self.total_params_df['bsa'] * 60
        
        print("Generated synthetic population data")
    
    def get_individual_parameters(self, subject_id=None, custom_params=None):
        """
        Generate parameters for individual PBPK modeling
        
        Args:
            subject_id: Specific subject ID to use (random if None)
            custom_params: Dictionary to override default parameters
        
        Returns:
            Dictionary of PBPK parameters
        """
        if self.total_params_df is None:
            raise ValueError("No population data available")
        
        # Select subject
        if subject_id is None:
            subject_data = self.total_params_df.sample(1).iloc[0]
        else:
            if subject_id <= len(self.total_params_df):
                subject_data = self.total_params_df.iloc[subject_id - 1]
            else:
                subject_data = self.total_params_df.sample(1).iloc[0]
        
        # Base parameters from population data
        params = {
            # Basic demographics
            'subject_id': int(subject_data['id']),
            'body_weight': float(subject_data['weight']),
            'height': float(subject_data['height']),
            'age': float(subject_data['age']),
            'sex': str(subject_data['sex']),
            'bsa': float(subject_data['bsa']),
            
            # Dose (fixed for demonstration)
            'dose': 100.0,
            
            # Absorption parameters (reasonable defaults)
            'ka': 1.0,      # Absorption rate constant (1/h)
            'fa': 0.9,      # Fraction absorbed
            'fg': 0.95,     # Gut availability
            'fh': 0.85,     # Hepatic availability
            'kd': 1.0,      # Dissolution rate constant
            'ts': 1.0,      # Stomach emptying time (h)
            'tlag': 0.1,    # Lag time (h)
            
            # Distribution volumes (scaled by body weight)
            'Vp': 0.043 * float(subject_data['weight']),    # Plasma volume (L)
            'Ve': 0.171 * float(subject_data['weight']),    # Extracellular volume (L)
            'Vt': 0.600 * float(subject_data['weight']),    # Tissue volume (L)
            
            # Protein binding and partitioning
            'fu': 0.1,      # Fraction unbound
            'BP': 1.0,      # Blood-plasma ratio
            
            # Tissue partition coefficients (reasonable defaults)
            'Kp_liver': 2.0,
            'Kp_kidney': 2.5,
            'Kp_brain': 0.5,
            'Kp_lung': 1.5,
            'Kp_heart': 1.8,
            'Kp_muscle': 1.2,
            'Kp_adipose': 0.8,
            'Kp_bone': 0.3,
            'Kp_skin': 1.0,
            'Kp_spleen': 2.0,
            
            # Clearance parameters (scaled by organ function)
            'CL_hepatic': max(0.5, 2.0 * float(subject_data.get('Liver_CYP3A4_activity', 1.0))),
            'CL_renal': max(0.3, 1.5 * (float(subject_data.get('Kidney_GFR', 120)) / 120)),
            'CL_biliary': 0.5,
            'CL_other': 0.2,
            
            # ADAM model parameters
            'use_adam': False,
            'solubility': 1.0,
            'permeability': 1e-4,
            'particle_size': 50.0,
            'density': 1.2,
            'molecular_weight': 300.0,
            'pKa': 7.4,
            'logP': 2.0
        }
        
        # Add organ-specific parameters if available
        organ_params = [
            'Liver_weight', 'Liver_volume', 'Liver_blood_flow',
            'Kidney_weight', 'Kidney_volume', 'Kidney_blood_flow',
            'Heart_weight', 'Heart_volume', 'Heart_blood_flow',
            'Lung_weight', 'Lung_volume', 'Lung_blood_flow'
        ]
        
        for param in organ_params:
            if param in subject_data:
                params[param] = float(subject_data[param])
        
        # Override with custom parameters
        if custom_params:
            params.update(custom_params)
        
        return params
    
    def get_population_parameters(self, n_subjects=100):
        """
        Generate parameters for population PBPK modeling
        
        Args:
            n_subjects: Number of subjects to simulate
            absorption_model: Type of absorption model ('first_order' or 'adam')
            distribution_model: Type of distribution model ('minimal_pbpk' or 'full_pbpk')
        
        Returns:
            List of parameter dictionaries
        """
        if self.total_params_df is None:
            raise ValueError("No population data available")
        
        # Select subjects
        if len(self.total_params_df) >= n_subjects:
            selected_subjects = self.total_params_df.sample(n_subjects)
        else:
            # If not enough unique subjects, sample with replacement
            selected_subjects = self.total_params_df.sample(n_subjects, replace=True)
        
        selected_subjects = selected_subjects.reset_index(drop=True)
        population_params = []
        
        for _, subject_data in selected_subjects.iterrows():
            # Get base parameters
            params = self.get_individual_parameters(subject_id=int(subject_data['id']))
            
            # Add biological variability to parameters
            
            # Add variability to key parameters
            np.random.seed(int(subject_data['id']))  # Reproducible randomness
            
            # Add biological variability
            params['ka'] *= np.random.lognormal(0, 0.3)  # 30% CV for absorption
            params['fa'] = np.random.beta(9, 1)  # Mean ~0.9
            params['fg'] = np.random.beta(9.5, 0.5)  # Mean ~0.95
            params['fh'] = np.random.beta(9, 1)  # Mean ~0.85
            params['fu'] = np.random.beta(2, 18)  # Mean ~0.1
            
            # Scale clearances by organ function
            params['CL_hepatic'] *= np.random.lognormal(0, 0.25)  # 25% CV for hepatic clearance
            params['CL_renal'] *= np.random.lognormal(0, 0.20)    # 20% CV for renal clearance
            
            # Tissue partition coefficients with variability
            for tissue in ['liver', 'kidney', 'brain', 'lung', 'heart', 'muscle', 'adipose']:
                key = f'Kp_{tissue}'
                if key in params:
                    params[key] *= np.random.lognormal(0, 0.2)  # 20% CV
            
            population_params.append(params)
        
        return population_params
    
    def save_parameters(self, params, filename):
        """Save parameters to JSON file"""
        # Convert numpy types to Python types for JSON serialization
        def convert_numpy(obj):
            if isinstance(obj, np.integer):
                return int(obj)
            elif isinstance(obj, np.floating):
                return float(obj)
            elif isinstance(obj, np.ndarray):
                return obj.tolist()
            return obj
        
        serializable_params = {k: convert_numpy(v) for k, v in params.items()}
        
        with open(filename, 'w') as f:
            json.dump(serializable_params, f, indent=2)
        
        print(f"Parameters saved to {filename}")
    
    def validate_parameters(self, params):
        """Validate generated parameters"""
        required_keys = ['dose', 'body_weight', 'ka', 'Vp', 'Ve', 'fu', 'CL_hepatic', 'CL_renal']
        
        missing_keys = [key for key in required_keys if key not in params]
        if missing_keys:
            raise ValueError(f"Missing required parameters: {missing_keys}")
        
        # Check for reasonable ranges
        checks = [
            ('dose', 0, 10000),
            ('body_weight', 30, 150),
            ('age', 18, 100),
            ('ka', 0.01, 10),
            ('Vp', 1, 10),
            ('fu', 0.01, 1.0),
            ('CL_hepatic', 0.1, 20),
            ('CL_renal', 0.1, 15)
        ]
        
        for param, min_val, max_val in checks:
            if param in params:
                val = float(params[param])
                if not (min_val <= val <= max_val):
                    print(f"Warning: {param} = {val} outside typical range [{min_val}, {max_val}]")


def main():
    """Test the parameter generator"""
    print("Testing Population Parameter Generator...")
    
    generator = PopulationParameterGenerator()
    
    # Test individual parameters
    print("\n1. Individual parameters:")
    individual_params = generator.get_individual_parameters(subject_id=1)
    print(f"   Subject ID: {individual_params['subject_id']}")
    print(f"   Body weight: {individual_params['body_weight']:.1f} kg")
    print(f"   Age: {individual_params['age']:.1f} years")
    print(f"   BSA: {individual_params['bsa']:.2f} m²")
    
    # Test population parameters
    print("\n2. Population parameters (5 subjects):")
    population_params = generator.get_population_parameters(n_subjects=5)
    for i, params in enumerate(population_params[:3]):
        print(f"   Subject {i+1}: {params['body_weight']:.1f} kg, {params['age']:.1f} yrs, {params['sex']}")
    
    # Save example parameters
    generator.save_parameters(individual_params, 'example_individual_params.json')
    generator.validate_parameters(individual_params)
    
    print("\nParameter generation completed!")


if __name__ == "__main__":
    main()